import os
import sys
sys.path.append('/share/project/duli/content_relation_ana/utils')
from utils import *
from collections import Counter
import random
import numpy as np
import pandas as pd
import os
import pdb

def tag_stat(id_ls, sample_ls):
    tag_num = []
    cate_num = []
    cate_ls = []
    for i in id_ls:
        tag_num.append(len(sample_ls[i]['label']['ability']))
        cate_num.append(len(sample_ls[i]['cate_ability']))
        cate_ls.extend(sample_ls[i]['cate_ability'])
    cate_counter = Counter(cate_ls)
    # import pdb
    # pdb.set_trace()
    tag_num = np.array(tag_num)
    cate_num = np.array(cate_num)
    return cate_counter, np.mean(tag_num), np.mean(cate_num), np.std(tag_num), np.std(cate_num)

def len_stat_dataset(sample_ls):
    dataset_dict = {}
    for sample in sample_ls:
        if sample['dataset'] not in dataset_dict.keys():
            dataset_dict[sample['dataset']] = []
        dataset_dict[sample['dataset']].append({'tot_len':sum([len(i['content']) for i in sample['content']]),
                                                'mean_instruction_len':np.mean([len(i['content']) for i in sample['content'] if i['role']=='user']),
                                                'mean_response_len':np.mean([len(i['content']) for i in sample['content'] if i['role']=='assistant']), 
                                                'num_turns':len(sample['content'])})
    res_dict = {}
    for dataset in dataset_dict.keys():
        res_tmp = pd.DataFrame(dataset_dict[dataset]).mean()
        res_dict[dataset] = res_tmp
    res_df = pd.DataFrame(res_dict)
    return res_df

def calculate_tag_and_cate_dataset_counts(samples):
    tag_ls = []
    cate_ls = []
    dataset_ls = []
    for sample in samples:
        tag_ls.extend(sample['label']['ability'])
        sample['cate_ability'] = [i.strip() for i in sample['cate_ability']]
        cate_ls.extend(sample['cate_ability'])
        dataset_ls.append(sample['dataset'])
    return tag_ls, cate_ls, dataset_ls


def tag_mean_loss(samples):
    ability_dict = {i:[] for i in set([tag for sample in samples for tag in sample['label']['ability']])}
    cate_dict = {i:[] for i in set([tag for sample in samples for tag in sample['cate_ability']])}
    
    for sample in samples:
        random.shuffle(sample['label']['ability'])
        random.shuffle(sample['cate_ability'])
        ability_dict[sample['label']['ability'][0]].append(sample['loss'])
        cate_dict[sample['cate_ability'][0]].append(sample['loss'])
    ability_mean_loss = {k:np.nanmean(v) for k,v in ability_dict.items()}
    ability_mean_loss = sorted([(k,v) for k,v in ability_mean_loss.items()], reverse=True, key=lambda x:x[1])
    cate_mean_loss = {k:np.nanmean(v) for k,v in cate_dict.items()}
    cate_mean_loss = sorted([(k,v) for k,v in cate_mean_loss.items()], reverse=True, key=lambda x:x[1])
    return ability_mean_loss, cate_mean_loss


dat_path = '/share/project/duli/content_relation_ana/subset_gen/loss_selection/tot_qwen1.5-7B'
sample_path = '/share/project/duli/content_relation_ana/subset_gen/sft_base_100w_v0513.jsonl'
sample_ls = load_jsonl(sample_path)

tot_sample_ls = []
file_ls = os.listdir(dat_path)
for file in file_ls:
    if '.jsonl' in file:
        try:
            #dat_tmp = load_jsonl(os.path.join(dat_path, file, 'cleaned_ability.jsonl'))
            dat_tmp = load_jsonl(os.path.join(dat_path, file))
            tot_sample_ls.extend(dat_tmp)
        except:
            pass

tot_sample_ls = sorted(tot_sample_ls, reverse=True, key=lambda x: float(x['loss']))
sample_dict = {(str(sample['id'])+sample['dataset']): sample for sample in sample_ls}

NUM = 200000
i = 0
addi_id_ls = []
while len(addi_id_ls) < NUM:
    key_tmp = str(tot_sample_ls[i]['id']) + tot_sample_ls[i]['dataset']
    try:
        _ = sample_dict[key_tmp]
    except:
        addi_id_ls.append(i)
    i += 1

addi_sample_ls = [tot_sample_ls[i] for i in addi_id_ls]

out_path = sample_path.replace('.jsonl', f'_highloss_{str(NUM)}.jsonl').replace('content_relation_ana/subset_gen/', 'content_relation_ana/subset_gen/loss_selection/')

# save_jsonl(sample_ls + addi_sample_ls, out_path)

len_stat = len_stat_dataset(tot_sample_ls)
# len_stat.to_csv('/share/project/duli/content_relation_ana/subset_gen/loss_selection/len_stat.csv')

ability_mean_loss, cate_mean_loss = tag_mean_loss(tot_sample_ls)

cate_ls = {}

for ith, sample in enumerate(tot_sample_ls):
    cate_tag_tmp = random.sample(sample['cate_ability'], 1)[0]
    if cate_tag_tmp not in cate_ls.keys():
        cate_ls[cate_tag_tmp] = []

    cate_ls[cate_tag_tmp].append(ith)

for cate in cate_ls.keys():
    cate_ls[cate] = sorted(cate_ls[cate], reverse=True, key=lambda x: tot_sample_ls[x]['loss'])

num_cate_ls = {k: int(len(cate_ls[k]) * NUM / len(tot_sample_ls)) +1 for k in  cate_ls.keys()}

num_cate_ls['自然语言处理与理解'] = 15000
num_cate_ls['信息处理与整合'] = 15000
num_cate_ls['数学能力'] = 30000
num_cate_ls['编程与软件开发'] = 30000
num_cate_ls['逻辑与推理'] = 30000


addi_sample_ls = []
for cate in cate_ls.keys():
    i = 0
    while i < num_cate_ls[cate]:
        sample_tmp = tot_sample_ls[cate_ls[cate][i]]
        key_tmp = str(sample_tmp['id']) + sample_tmp['dataset']
        try:
            _ = sample_dict[key_tmp]
        except:
            addi_sample_ls.append(tot_sample_ls[i])
        i += 1

addi_logi_sample = load_jsonl('/share/project/duli/sft_dataset_collections/logi-cot/reformat.jsonl')
addi_logi_sample = random.sample(addi_logi_sample, 10000)

out_path = sample_path.replace('.jsonl', f'_highloss_cate_{str(NUM)}.jsonl').replace('content_relation_ana/subset_gen/', 'content_relation_ana/subset_gen/loss_selection/')
save_jsonl(sample_ls + addi_sample_ls + addi_logi_sample, out_path)
pdb.set_trace()

with open('/share/project/duli/content_relation_ana/subset_gen/loss_selection/ability_mean_loss.txt', 'w') as f:
    for i in ability_mean_loss:
        f.writelines(str(i) + '\n')    
    
with open('/share/project/duli//content_relation_ana/subset_gen/loss_selection/cate_mean_loss.txt', 'w') as f:
    for i in cate_mean_loss:
        f.writelines(str(i) + '\n')    
    



